{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "from transformers import Trainer, TrainerCallback, TrainingArguments\n",
    "from arguments import SFTArguments\n",
    "from datasets import load_dataset\n",
    "from qwen.modeling_qwen import QWenLMHeadModel\n",
    "from qwen.tokenization_qwen import QWenTokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sft_args = SFTArguments()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PROMPT_DICT = {\n",
    "    \"prompt_input\": (\"你是一个助手 \" \"用户: {instruction} {input} 回答: \"),\n",
    "    \"prompt_no_input\": (\"你是一个助手 \" \"用户: {instruction}  回答: \"),\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = load_dataset(\n",
    "    path=\"parquet\", data_files=sft_args.SFT_FILES, split=\"train\", keep_in_memory=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = QWenTokenizer.from_pretrained(sft_args.tokenizer_dir)\n",
    "print(f\"vicab size: {len(tokenizer)}\")\n",
    "tokenizer.pad_token_id = tokenizer.im_end_id\n",
    "map_dtype = np.uint16 if len(tokenizer) < 65535 else np.uint32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_example(example):\n",
    "    prompt_input, prompt_no_input = (\n",
    "        PROMPT_DICT[\"prompt_input\"],\n",
    "        PROMPT_DICT[\"prompt_no_input\"],\n",
    "    )\n",
    "    if example.get(\"input\"):\n",
    "        target = example[\"output\"] + \"<|im_end|>\"\n",
    "        context = prompt_input.format_map(\n",
    "            dict(instruction=example[\"instruction\"], input=example[\"input\"])\n",
    "        )\n",
    "\n",
    "        example[\"context\"] = context\n",
    "        example[\"target\"] = target\n",
    "    else:\n",
    "        target = example[\"output\"] + \"<|im_end|>\"\n",
    "        context = prompt_no_input.format_map(dict(instruction=example[\"instruction\"]))\n",
    "\n",
    "        example[\"context\"] = context\n",
    "        example[\"target\"] = target\n",
    "    return example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess(example):\n",
    "    prompt = example[\"context\"]\n",
    "    target = example[\"target\"]\n",
    "    input_ids = tokenizer(\n",
    "        prompt + target,\n",
    "        return_tensors=\"pt\",\n",
    "        padding=\"longest\",\n",
    "        max_length=512,\n",
    "        truncation=True,\n",
    "    )\n",
    "    seq_ids = tokenizer(\n",
    "        prompt,\n",
    "        return_tensors=\"pt\",\n",
    "        padding=\"longest\",\n",
    "        max_length=512,\n",
    "        truncation=True,\n",
    "    )\n",
    "    input_ids_len = seq_ids.input_ids.ne(tokenizer.pad_token_id).sum().item()\n",
    "\n",
    "    return {\"input_ids\": input_ids.input_ids[0], \"seq_len\": input_ids_len}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_datasets = dataset.map(\n",
    "    function=format_example, num_proc=32, keep_in_memory=False\n",
    ")\n",
    "print(\"1\")\n",
    "print(tokenized_datasets)\n",
    "tokenized_datasets = tokenized_datasets.map(\n",
    "    function=preprocess, num_proc=32, keep_in_memory=False\n",
    ").shuffle(23333)\n",
    "print(\"2\")\n",
    "print(tokenized_datasets)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def data_collator(fetures):\n",
    "    len_ids = [len(feture[\"input_ids\"]) for feture in fetures]\n",
    "    longest = max(len_ids) + 1\n",
    "    input_ids = []\n",
    "    attention_mask_list = []\n",
    "    postion_ids_list = []\n",
    "    labels_list = []\n",
    "    for ids_l, feture in sorted(zip(len_ids, fetures), key=lambda x: -x[0]):\n",
    "        ids = feture[\"input_ids\"]\n",
    "        seq_len = feture[\"seq_len\"]\n",
    "        labels = [-100] * seq_len + ids[seq_len:] + [-100] * (longest - ids_l)\n",
    "        ids = ids + [tokenizer.im_end_id] * (longest - ids_l)\n",
    "        _ids = torch.LongTensor(ids)\n",
    "        labels_list.append(torch.LongTensor(labels))\n",
    "        input_ids.append(_ids)\n",
    "    input_ids = torch.stack(input_ids)\n",
    "    labels = torch.stack(labels_list)\n",
    "\n",
    "    return {\"input_ids\": input_ids, \"labels\": labels}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = QWenLMHeadModel.from_pretrained(sft_args.sft_from_checkpoint_file)\n",
    "\n",
    "model_size = sum(t.numel() for t in model.parameters())\n",
    "print(f\"Qwen size: {model_size / 1000**2:.2f}M parameters\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EmptyCudaCacheCallback(TrainerCallback):\n",
    "    log_cnt = 0\n",
    "\n",
    "    def on_log(self, args, state, control, logs=None, **kwargs):\n",
    "        self.log_cnt += 1\n",
    "        if self.log_cnt % 5 == 0:\n",
    "            torch.cuda.empty_cache()\n",
    "\n",
    "\n",
    "empty_cuda_cahce = EmptyCudaCacheCallback()\n",
    "\n",
    "\n",
    "my_datasets = tokenized_datasets.train_test_split(test_size=4096)\n",
    "print(\"m\")\n",
    "print(my_datasets)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "args = TrainingArguments(\n",
    "    output_dir=sft_args.model_save_dir,\n",
    "    per_device_train_batch_size=32,\n",
    "    gradient_accumulation_steps=2,\n",
    "    num_train_epochs=3,\n",
    "    weight_decay=0.1,\n",
    "    warmup_steps=0,\n",
    "    learning_rate=6e-5,\n",
    "    ddp_find_unused_parameters=False,\n",
    "    evaluation_strategy=\"steps\",\n",
    "    eval_steps=500,\n",
    "    save_steps=500,\n",
    "    save_total_limit=3,\n",
    "    report_to=\"tensorboard\",\n",
    "    optim=\"adamw_torch\",\n",
    "    remove_unused_columns=False,\n",
    "    lr_scheduler_type=\"cosine\",\n",
    "    bf16=True,\n",
    "    logging_steps=10,\n",
    "    log_level=\"info\",\n",
    "    logging_first_step=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    tokenizer=tokenizer,\n",
    "    args=args,\n",
    "    data_collator=data_collator,\n",
    "    train_dataset=my_datasets[\"train\"],\n",
    "    eval_dataset=my_datasets[\"test\"],\n",
    "    callbacks=[empty_cuda_cahce],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_results = trainer.evaluate()\n",
    "print(f\"Perplexity: {np.exp(eval_results['eval_loss']):.2f}\")\n",
    "loss_log = pd.DataFrame(trainer.state.log_history)\n",
    "trainer.save_model(sft_args.model_save_dir)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "d2l",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.21"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
